import numpy as np
import quimb as qu
import quimb.tensor as qtn
from quimb.tensor import Tensor, TensorNetwork

I, X, Y, Z = np.eye(2), np.array([[0.0, 1], [1, 0]]), np.array([[0.0, -1j], [1j, 0]]), np.array([[1.0, 0], [0, -1]])
Z_mea_o1 = np.array([0, 0, 0, np.sqrt(2)])
meas_map = {
	'z+': np.array([1, 0]), 'z-': np.array([0, 1]),
}
"""Setting"""
cached = True
N_qubit, N_set, N_shot, step = 10, 1, 2000, 5
J, h, dt = 0.5236, 1.0, 0.05
max_bond, cutoff = 400, 1e-10

"""simulation"""
# probs = [1 / 3, 1 / 3, 1 / 3]
# px, py, pz = probs
# dual_o2 = np.array([
# 	[1, 1, 1, 1, 1, 1],
# 	[1 / px, -1 / px, 0, 0, 0, 0],
# 	[0, 0, 1 / py, -1 / py, 0, 0],
# 	[0, 0, 0, 0, 1 / pz, -1 / pz]
# ]) / np.sqrt(2)
dual_o2 = np.array([
	[1, 1],
	[0, 0],
	[0, 0],
	[1, -1]
]) / np.sqrt(2)
site_tags = list(map("I{}".format, range(N_qubit)))

import pathlib

"""Error Mitigation"""
import tqdm
import json

Zp_o1 = np.array([1, 0, 0, 1]) / np.sqrt(2)
rho0_mps = TensorNetwork([Tensor(Zp_o1, inds=(f"b{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)], virtual=True)

Z_meas_mps = TensorNetwork([Tensor(Z_mea_o1, inds=(f"k{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)], virtual=True)
dual_mpo = TensorNetwork([Tensor(dual_o2, inds=(f"b{_}", f"c{_}"), tags=[f"I{_}"]) for _ in range(N_qubit)],
						 virtual=True)

# steps = [0, 1, 2, 3, 4, 5, 6, 7]
steps = np.arange(step + 1)
mean_val, std_val = [1.0], [0.]
ideal_mean_vals, noisy_mean_vals = [1.0], [1.0]
meas_vals = []
for i in steps[1:]:
	"""Data Load"""
	data_filename = f"data/Outcomes_reduced, N_qubit={N_qubit}, N_set={N_set}, N_shot={N_shot}, step={i}, max_bond={max_bond}.json"
	# data_filename = f"data/Outcomes, N_qubit={N_qubit}, N_set={N_set}, N_shot={N_shot}, step={i}, max_bond={max_bond}.json"
	file = pathlib.Path(data_filename)
	if not file.exists(): raise FileNotFoundError(f"Data File '{data_filename}' not Exist!")
	with open(file, 'r') as f:
		print(f"Loading Measurement data from File '{data_filename}' ...")
		data_dict = json.load(f)
		print(f"Loading done.")

	datas = data_dict['datas']
	"""TEM Load"""
	TEM_trotter_step_filename = f"cache/TEM_trotter_step, step={i}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	if not pathlib.Path(TEM_trotter_step_filename).exists(): raise FileNotFoundError
	TEM_trotter_step = qu.load_from_disk(TEM_trotter_step_filename)
	tn = dual_mpo | TEM_trotter_step | Z_meas_mps
	qtn.tensor_network_1d_compress(tn, max_bond=4 * max_bond, cutoff=cutoff, site_tags=site_tags,
								   method='direct', inplace=True)

	estimated_vals = []
	for j in range(N_set):
		data = datas[j]
		meas_setting, outcomes = data['meas'], data['outcomes']
		pbar = tqdm.trange(N_shot, desc=f'step {i}, TEM {j + 1}/{N_set}', leave=True)
		for k in pbar:
			outcome = outcomes[k]
			mea_mps = TensorNetwork(
					[Tensor(meas_map[outcome[l]], inds=(f"c{l}",), tags=[f"I{l}"]) for l in range(N_qubit)],
					virtual=True)
			estimated_vals.append(mea_mps @ tn)
		pbar.update(1)

	estimated_vals = np.asarray(estimated_vals)
	meas_vals.append(estimated_vals)

	mean, std = np.mean(estimated_vals), np.std(estimated_vals)
	mean_val.append(mean)
	std_val.append(std)

	print(f"step {i}, mean={mean:.3f}, std={std:.3f}")

	trotter_step_filename = f"cache/trotter_step, step={i}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	trotter_step = qu.load_from_disk(trotter_step_filename)
	noisy_trotter_step_filename = f"cache/noisy_trotter_step, step={i}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
	noisy_trotter_step = qu.load_from_disk(noisy_trotter_step_filename)

	ideal_mean_vals.append((rho0_mps | trotter_step | Z_meas_mps).contract())
	noisy_mean_vals.append((rho0_mps | noisy_trotter_step | Z_meas_mps).contract())

import matplotlib.pyplot as plt

plt.boxplot(meas_vals, showmeans=True, showbox=False, showfliers=False)

plt.title(f'max_bond={max_bond}, cutoff={cutoff:.1e}', fontsize=18)
plt.plot(steps, ideal_mean_vals, 'k--', label='exact')
plt.plot(steps, noisy_mean_vals, color='red', label='noisy')
plt.plot(steps, mean_val, color='green', label='TEM')

plt.xlim([0, step])
plt.ylim([-1.5, 1.5])
plt.xlabel('steps', fontsize=14)
plt.ylabel(r'$\langle Z^{\otimes 10}\rangle$', fontsize=14)
plt.legend()
plt.savefig('TEM.png')
plt.show()

plt.plot(steps, std_val, 'g-', marker='o')
plt.xlim([0, step])
plt.ylim([0., max(std_val) + 1])
plt.xlabel('steps', fontsize=14)
plt.ylabel('$\Delta Z^{\otimes 10}$', fontsize=14)
plt.show()
